# -*- encoding: utf-8 -*-
"""
@File    :   dataset.py
@Contact :   thgpddl@163.com

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2022/5/18 23:45   thgpddl      1.0         None
"""
import numpy as np
import pandas as pd
import torch
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torch.utils.data import Dataset
import os




def get_ckplusdataloaders(path="D:\WorkSpace\zh_clearml_pt\dataset\ck+\ck+2", bs=64, num_workers=0, augment=True):
    """

    @param path: csv文件路径
    @param bs: 训练集batchsize
    @param ebs: 验证集和测试集batchsize
    @param num_workers:
    @param augment: 训练集是否进行数据增强
    @return:
    """



    mu, st = 0, 1

    test_transform = transforms.Compose([
        transforms.Grayscale(),
        transforms.Resize(48),
        transforms.TenCrop(40),
        transforms.Lambda(lambda crops: torch.stack(
            [transforms.ToTensor()(crop) for crop in crops])),
        transforms.Lambda(lambda tensors: torch.stack(
            [transforms.Normalize(mean=(mu,), std=(st,))(t) for t in tensors]))
    ])
    if augment:
        train_transform = transforms.Compose([
            transforms.Grayscale(),
            transforms.Resize(48),
            transforms.RandomResizedCrop(48, scale=(0.8, 1.2)),
            transforms.RandomApply([transforms.ColorJitter(
                brightness=0.5, contrast=0.5, saturation=0.5)], p=0.5),
            transforms.RandomApply(
                [transforms.RandomAffine(0, translate=(0.2, 0.2))], p=0.5),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([transforms.RandomRotation(10)], p=0.5),
            transforms.TenCrop(40),
            transforms.Lambda(lambda crops: torch.stack(
                [transforms.ToTensor()(crop) for crop in crops])),
            transforms.Lambda(lambda tensors: torch.stack(
                [transforms.Normalize(mean=(mu,), std=(st,))(t) for t in tensors])),
            transforms.Lambda(lambda tensors: torch.stack(
                [transforms.RandomErasing()(t) for t in tensors])),
        ])
    else:
        train_transform =test_transform

    train=ImageFolder(root=os.path.join(path,'train'),transform=train_transform)
    val=ImageFolder(root=os.path.join(path,'test'),transform=test_transform)

    trainloader = DataLoader(train, batch_size=bs, shuffle=True, num_workers=num_workers)
    valloader = DataLoader(val, batch_size=bs, shuffle=True, num_workers=num_workers)

    return trainloader, valloader, valloader


if __name__ == "__main__":
    train_loader, val_loader, test_loader = get_ckplusdataloaders(bs=4,
                                                               num_workers=0,
                                                               augment=True)
    data=next(iter(train_loader))
    images,targets=data  # bs, ten,1,h,w
    images=images.view(-1,1,images.size(-2),images.size(-1))
    from visual.visual_tensor import show_tensor
    show_tensor(tensor=images,nrow=4,save_path="train_loader.jpg")
